import numpy as np
import argparse
import re


class DatasetViewer:
	def __init__(self, data_path):
		"""初始化数据集查看器"""
		self.data_path = data_path
		self.data = None
		self.word2ix = None
		self.ix2word = None
		self.poems = None

	def load_data(self):
		"""加载数据集"""
		try:
			print(f"正在加载数据集: {self.data_path}")
			self.data = np.load(self.data_path, allow_pickle=True)

			# 获取词汇表和诗词数据
			self.word2ix = self.data['word2ix'].item()
			self.ix2word = self.data['ix2word'].item()
			self.poems = self.data['data']

			print(f"数据集加载成功: {len(self.poems)} 首诗, 词汇表大小: {len(self.word2ix)}")
			print(f"特殊标记:")
			for token in ['<START>', '<EOP>', '</s>', '<UNK>']:
				if token in self.word2ix:
					print(f"  {token}: 索引 {self.word2ix[token]}")
				else:
					print(f"  {token}: 未在词汇表中找到")

		except Exception as e:
			print(f"加载数据集失败: {e}")
			return False

		return True

	def view_sample(self, index=None):
		"""查看单个样本"""
		if self.poems is None:
			print("请先加载数据集")
			return

		# 如果未指定索引，随机选择一个样本
		if index is None:
			index = np.random.randint(0, len(self.poems))

		# 确保索引有效
		if index < 0 or index >= len(self.poems):
			print(f"索引超出范围 (0-{len(self.poems) - 1})")
			return

		poem_indices = self.poems[index]

		# 转换为文本
		poem_text = ''.join([self.ix2word[idx] for idx in poem_indices if idx in self.ix2word])

		print(f"\n样本 #{index}:")
		print(f"  索引序列长度: {len(poem_indices)}")
		print(f"  文本内容: {poem_text}")

		# 尝试检测诗的类型（五言/七言）
		if len(poem_text) > 10:
			# 假设逗号和句号是行分隔符
			lines = poem_text.replace('，', '\n').replace('。', '\n').split('\n')
			lines = [line.strip() for line in lines if line.strip()]

			if len(lines) > 1:
				line_lengths = [len(line) for line in lines]
				avg_length = sum(line_lengths) / len(line_lengths)

				print(f"  可能的诗类型: {int(round(avg_length))}言诗 ({len(lines)}行)")

		return poem_indices, poem_text

	def analyze_data(self):
		"""分析数据集统计信息"""
		if self.poems is None:
			print("请先加载数据集")
			return

		print("\n===== 数据集统计信息 =====")

		# 诗词长度分布
		poem_lengths = [len(p) for p in self.poems]
		avg_length = sum(poem_lengths) / len(poem_lengths)

		print(f"诗词长度分布:")
		print(f"  平均长度: {avg_length:.2f} 个字符")
		print(f"  最小长度: {min(poem_lengths)} 个字符")
		print(f"  最大长度: {max(poem_lengths)} 个字符")

		# 常见字符统计
		char_counts = {}
		for poem in self.poems:
			for idx in poem:
				char = self.ix2word.get(idx, '<UNK>')
				char_counts[char] = char_counts.get(char, 0) + 1

		# 按频率排序
		sorted_chars = sorted(char_counts.items(), key=lambda x: x[1], reverse=True)

		print("\n最常见的10个字符:")
		for i, (char, count) in enumerate(sorted_chars[:10]):
			print(f"  {i + 1}. '{char}': {count} 次")

		print("\n最不常见的10个字符:")
		for i, (char, count) in enumerate(sorted_chars[-10:]):
			print(f"  {i + 1}. '{char}': {count} 次")

		# 特殊标记出现频率
		print("\n特殊标记出现频率:")
		for token in ['<START>', '<EOP>', '</s>', '<UNK>']:
			if token in self.word2ix:
				token_idx = self.word2ix[token]
				count = char_counts.get(token, 0)
				print(f"  {token}: {count} 次")

	def analyze_poem_types(self):
		"""分析诗的类型（五言/七言）"""
		if self.poems is None:
			print("请先加载数据集")
			return

		five_poem_count = 0
		seven_poem_count = 0
		other_poem_count = 0

		for poem_indices in self.poems:
			poem_text = ''.join([self.ix2word.get(idx, '') for idx in poem_indices])
			poem_text = re.sub(r'[^\w\s]', '', poem_text)  # 去除非文字字符

			# 检查是否为五言诗（每句5字）
			is_five = True
			for i in range(0, len(poem_text), 5):
				line = poem_text[i:i + 5]
				if len(line) != 5:
					is_five = False
					break
			if is_five:
				five_poem_count += 1
				continue

			# 检查是否为七言诗（每句7字）
			is_seven = True
			for i in range(0, len(poem_text), 7):
				line = poem_text[i:i + 7]
				if len(line) != 7:
					is_seven = False
					break
			if is_seven:
				seven_poem_count += 1
				continue

			other_poem_count += 1

		print(f"\n===== 诗类型分析 =====")
		print(f"五言诗数量: {five_poem_count}")
		print(f"七言诗数量: {seven_poem_count}")
		print(f"其他格式数量: {other_poem_count}")


def main():
	parser = argparse.ArgumentParser(description='古诗数据集查看器')
	parser.add_argument('--data_path', type=str, default='data/tang.npz', help='数据集路径')
	parser.add_argument('--sample_index', type=int, default=None, help='要查看的样本索引，随机选择则不指定')
	parser.add_argument('--analyze', action='store_true', help='是否分析整个数据集')
	parser.add_argument('--poem_type_analyze', action='store_true', help='分析诗的类型（五言/七言）')

	args = parser.parse_args()

	viewer = DatasetViewer(args.data_path)

	if viewer.load_data():
		if args.sample_index is not None:
			viewer.view_sample(args.sample_index)
		else:
			viewer.view_sample()

		if args.analyze:
			viewer.analyze_data()

		if args.poem_type_analyze:
			viewer.analyze_poem_types()


if __name__ == "__main__":
	main()